{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "from sklearn.preprocessing import OneHotEncoder\n",
    "import numpy as np\n",
    "import pandas as pd \n",
    "\n",
    "\n",
    "#获取特征-域字典。格式：{feature_index1:field_index1,......}\n",
    "def onehot_get_dict(train_x,test_x):\n",
    "    #未onehot编码前，原特征数，即域的个数\n",
    "    m = len(train_x.columns)\n",
    "    \n",
    "    #onehot\n",
    "    x_all = np.concatenate([train_x, test_x])\n",
    "    enc = OneHotEncoder(sparse = False)\n",
    "    enc.fit(x_all)\n",
    "    x_all = enc.transform(x_all)\n",
    "    train_x = enc.transform(train_x)\n",
    "    test_x = enc.transform(test_x)\n",
    "    \n",
    "    #one-hot后的特征数\n",
    "    n = x_all.shape[1]\n",
    "    del x_all\n",
    "    \n",
    "    #m_list为one-hot后，每个特征的取值数目(即每个原特征分解成多少个子特征)\n",
    "    #格式为[n1,n2,n3....],size为m原始特征数。即第一个特征(域)one-hot编码后,有n1种不同取值.....\n",
    "    m_list = list(enc.n_values_)\n",
    "    \n",
    "    dic = dict()\n",
    "    h = 0\n",
    "    #循环每个域\n",
    "    for i in range(m):\n",
    "        #循环每个域中的特征数,更新字典\n",
    "        for j in range(m_list[0]):\n",
    "            dic[h] = i\n",
    "            h = h+1\n",
    "        m_list.pop(0)\n",
    "    \n",
    "    #稀疏矩阵转换成dataFrame\n",
    "    \n",
    "    train_x = pd.DataFrame(train_x)\n",
    "    test_x = pd.DataFrame(test_x)\n",
    "    \n",
    "    return train_x,test_x,n,m,dic"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
